import os

import pandas as pd
from tensorflow.keras.preprocessing import image
from tensorflow.keras.preprocessing.image import ImageDataGenerator


def adjust_data(img, mask):
    img = img / 255
    mask = mask / 255
    mask[mask > 0.5] = 1
    mask[mask <= 0.5] = 0
    return img, mask


def train_data_generator(arguments, params):
    target_size = (params['H'], params['W'])
    random_seed = params['random_seed']
    data_path = arguments.data_path

    image_data_gen = ImageDataGenerator(**params['data_augmentation'])
    mask_data_gen = ImageDataGenerator(**params['data_augmentation'])
    image_generator = image_data_gen.flow_from_directory(
        data_path,
        classes=['image'],
        class_mode=None,
        color_mode='grayscale',
        target_size=target_size,
        batch_size=1,
        seed=random_seed)
    mask_generator = mask_data_gen.flow_from_directory(
        data_path,
        classes=['label'],
        class_mode=None,
        color_mode='grayscale',
        target_size=target_size,
        batch_size=1,
        seed=random_seed)
    return zip(image_generator, mask_generator)


def entrypoint(arguments, params):
    # TODO: check with output naming.
    train_generator = train_data_generator(arguments, params)
    cont_list = []

    i = 0
    save_dir = arguments.output_path
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    if not os.path.exists(os.path.join(save_dir, f'image')):
        os.mkdir(os.path.join(save_dir, f'image'))
    if not os.path.exists(os.path.join(save_dir, f'label')):
        os.mkdir(os.path.join(save_dir, f'label'))
    for (img, mask) in train_generator:
        img, mask = adjust_data(img, mask)
        image_name = os.path.join(save_dir, f'image/{i}.png')
        mask_name = os.path.join(save_dir, f'label/{i}.png')
        img = img[0, :, :, :]
        mask = mask[0, :, :, :]
        image.save_img(image_name, img)
        image.save_img(mask_name, mask)
        cont_list.append({'image_path': f'image/{i}.png', 'mask_path': f'label/{i}.png'})
        i = i + 1
        if i == 200:
            break
    df = pd.DataFrame(cont_list, columns=['image_path', 'mask_path'])
    df.to_csv(f"{save_dir}/train.csv", index=False)
